# -*- coding: utf-8 -*-
"""The plan module related tests."""
from unittest import IsolatedAsyncioTestCase

from agentscope.plan import SubTask, Plan, PlanNotebook


class PlanTest(IsolatedAsyncioTestCase):
    """Test the plan module."""

    async def asyncSetUp(self) -> None:
        """Set up the test case."""
        self.subtask1 = SubTask(
            name="Task 1",
            description="Description 1",
            expected_outcome="Expected outcome 1",
            state="done",
        )

        self.subtask2 = SubTask(
            name="Task 2",
            description="Description 2",
            expected_outcome="Expected outcome 2",
            state="in_progress",
        )

        self.subtask3 = SubTask(
            name="Task 3",
            description="Description 3",
            expected_outcome="Expected outcome 3",
            state="todo",
        )

        self.plan = Plan(
            name="Create website",
            description="Create a personal portfolio website.",
            expected_outcome="A new personal portfolio website is created on "
            "GitHub.",
            subtasks=[self.subtask1, self.subtask2, self.subtask3],
        )

    async def test_plan_model(self) -> None:
        """Test the models used in plan module."""

        self.assertEqual(
            self.subtask1.to_markdown(detailed=True),
            f"""- [x] Task 1
\t- Created At: {self.subtask1.created_at}
\t- Description: Description 1
\t- Expected Outcome: Expected outcome 1
\t- State: done
\t- Finished At: None
\t- Actual Outcome: None""",
        )

        self.assertEqual(
            self.subtask1.to_markdown(detailed=False),
            """- [x] Task 1""",
        )

        self.assertEqual(
            self.plan.to_markdown(detailed=True),
            f"""# Create website
**Description**: Create a personal portfolio website.
**Expected Outcome**: A new personal portfolio website is created on GitHub.
**State**: todo
**Created At**: {self.plan.created_at}
## Subtasks
- [x] Task 1
\t- Created At: {self.subtask1.created_at}
\t- Description: Description 1
\t- Expected Outcome: Expected outcome 1
\t- State: done
\t- Finished At: None
\t- Actual Outcome: None
- [ ] [WIP]Task 2
\t- Created At: {self.subtask2.created_at}
\t- Description: Description 2
\t- Expected Outcome: Expected outcome 2
\t- State: in_progress
- [ ] Task 3
\t- Created At: {self.subtask3.created_at}
\t- Description: Description 3
\t- Expected Outcome: Expected outcome 3
\t- State: todo""",
        )

        self.assertEqual(
            self.plan.to_markdown(detailed=True),
            f"""# Create website
**Description**: Create a personal portfolio website.
**Expected Outcome**: A new personal portfolio website is created on GitHub.
**State**: todo
**Created At**: {self.plan.created_at}
## Subtasks
- [x] Task 1
\t- Created At: {self.subtask1.created_at}
\t- Description: Description 1
\t- Expected Outcome: Expected outcome 1
\t- State: done
\t- Finished At: None
\t- Actual Outcome: None
- [ ] [WIP]Task 2
\t- Created At: {self.subtask2.created_at}
\t- Description: Description 2
\t- Expected Outcome: Expected outcome 2
\t- State: in_progress
- [ ] Task 3
\t- Created At: {self.subtask3.created_at}
\t- Description: Description 3
\t- Expected Outcome: Expected outcome 3
\t- State: todo""",
        )

    async def test_plan_subtasks(self) -> None:
        """Test the plan and subtask models."""
        plan_notebook = PlanNotebook()

        self.assertListEqual(
            [_.__name__ for _ in plan_notebook.list_tools()],
            [
                "view_subtasks",
                "update_subtask_state",
                "finish_subtask",
                "create_plan",
                "revise_current_plan",
                "finish_plan",
                "view_historical_plans",
                "recover_historical_plan",
            ],
        )

        plan_hint = await plan_notebook.get_current_hint()
        self.assertEqual(
            plan_hint.get_text_content(),
            "<system-hint>If the user's query is complex (e.g. "
            "programming a website, game or app), or requires a long chain of "
            "steps to complete (e.g. conduct research on a certain topic from "
            "different sources), you NEED to create a plan first by calling "
            "'create_plan'. Otherwise, you can directly execute the user's "
            "query without planning.</system-hint>",
        )

        res = await plan_notebook.create_plan(
            name="Example Plan",
            description="Example Description",
            expected_outcome="Example Expected Outcome",
            subtasks=[self.subtask1, self.subtask2, self.subtask3],
        )
        self.assertEqual(
            res.content[0]["text"],
            "Plan 'Example Plan' created successfully.",
        )

        res = await plan_notebook.view_subtasks([3])
        self.assertEqual(
            res.content[0]["text"],
            "Invalid subtask_idx '[3]'. Must be between 0 and 2.",
        )
        res = await plan_notebook.view_subtasks([0, 2])
        self.assertEqual(
            res.content[0]["text"],
            f"""Subtask at index 0:
```
- [x] Task 1
\t- Created At: {self.subtask1.created_at}
\t- Description: Description 1
\t- Expected Outcome: Expected outcome 1
\t- State: done
\t- Finished At: None
\t- Actual Outcome: None
```

Subtask at index 2:
```
- [ ] Task 3
\t- Created At: {self.subtask3.created_at}
\t- Description: Description 3
\t- Expected Outcome: Expected outcome 3
\t- State: todo
```
""",
        )

        await plan_notebook.revise_current_plan(
            1,
            action="add",
            subtask=SubTask(
                name="Task 11",
                description="Description 11",
                expected_outcome="Expected outcome 11",
            ),
        )
        self.assertEqual(
            plan_notebook.current_plan.subtasks[1].name,
            "Task 11",
        )
        self.assertEqual(
            len(plan_notebook.current_plan.subtasks),
            4,
        )

        res = await plan_notebook.revise_current_plan(
            1,
            "delete",
        )
        self.assertEqual(
            res.content[0]["text"],
            "Subtask (named 'Task 11') at index 1 is deleted successfully.",
        )
        self.assertEqual(
            len(plan_notebook.current_plan.subtasks),
            3,
        )

        res = await plan_notebook.revise_current_plan(
            1,
            "revise",
            subtask=SubTask(
                name="Task 22",
                description="Description 22",
                expected_outcome="Expected outcome 22",
            ),
        )
        self.assertEqual(
            res.content[0]["text"],
            "Subtask at index 1 is revised successfully.",
        )
        self.assertEqual(
            plan_notebook.current_plan.subtasks[1].name,
            "Task 22",
        )
        self.assertEqual(
            len(plan_notebook.current_plan.subtasks),
            3,
        )

        res = await plan_notebook.update_subtask_state(
            2,
            "in_progress",
        )
        self.assertEqual(
            res.content[0]["text"],
            "Subtask (at index 1) named 'Task 22' is not done yet. "
            "You should finish the previous subtasks first.",
        )

        await plan_notebook.update_subtask_state(0, "in_progress")
        res = await plan_notebook.update_subtask_state(
            1,
            "in_progress",
        )
        self.assertEqual(
            res.content[0]["text"],
            "Subtask (at index 0) named 'Task 1' is not done yet. You "
            "should finish the previous subtasks first.",
        )

        res = await plan_notebook.finish_subtask(
            0,
            "Fake outcome for task 1",
        )
        self.assertEqual(
            res.content[0]["text"],
            "Subtask (at index 0) named 'Task 1' is marked as done "
            "successfully. The next subtask named 'Task 22' is activated.",
        )
        self.assertEqual(
            plan_notebook.current_plan.subtasks[1].state,
            "in_progress",
        )
